import os
import requests
import json
import concurrent.futures
import base64
import paddleclas
import time
import argparse
import traceback
import torch
import paddle
from paddleocr import PaddleOCR
from modelscope.pipelines import pipeline 
from modelscope.utils.constant import Tasks 
from flask import Flask, request, current_app
from shapely.geometry import Polygon
from modelscope.models.cv.ocr_detection.utils import boxes_from_bitmap
from service.ocr import pre_process_images, angle_images, read_images_ocr_words, read_images_tables
from service.barcode import read_images_barcodes
from service.image import compare_images

SERVICE_TYPE = "SERVICE_TYPE"
MODEL_PADDLE_OCR = "MODEL_PADDLE_OCR"
MODEL_OCR_RECOGNITION = "MODEL_OCR_RECOGNITION"
MODEL_OCR_DETERTION = "MODEL_OCR_DETERTION"
MODEL_TABLE_RECOGNITION = "MODEL_TABLE_RECOGNITION"
MODEL_IMAGE_ORIENTATION_PREDICTOR = "MODEL_IMAGE_ORIENTATION_PREDICTOR"
executor = concurrent.futures.ThreadPoolExecutor(max_workers=3)

def initialize_app(app, service_type):
    paddle_ocr = None
    ocr_detection = None
    ocr_recognition = None
    table_recognition = None
    image_orientation_predictor = None

    start_time = time.time()
    print(f"initialize app start, service = {service_type}")
    image_orientation_predictor = paddleclas.PaddleClas(model_name="text_image_orientation")
    if service_type == 'PF':
        paddle_ocr = PaddleOCR(use_gpu=True, show_log=False, det_db_unclip_ratio=1.2, use_space_char=True, det_db_box_thresh=0.3, use_dilation=True, det_limit_side_len=1280, use_angle_cls=False, lang='ch', det_model_dir='/root/.paddleocr/paddleocr_det_server_infer', rec_model_dir='/root/.paddleocr/paddleocr_rec_server_infer', cls_model_dir='/root/.paddleocr/paddleocr_cls_infer')
        ocr_detection = pipeline(Tasks.ocr_detection, model='/root/.modelscope/damo/cv_resnet18_ocr-detection-db-line-level_damo')
        ocr_recognition = pipeline(Tasks.ocr_recognition, model='/root/.modelscope/damo/cv_convnextTiny_ocr-recognition-document_damo')
        # 取消对表格识别的支持，因为不准
        # table_recognition = pipeline(Tasks.table_recognition, model='/root/.modelscope/damo/cv_dla34_table-structure-recognition_cycle-centernet')
    else:
        paddle_ocr = PaddleOCR(use_gpu=True, show_log=False, det_db_unclip_ratio=1.2, use_space_char=True, det_db_box_thresh=0.3, use_dilation=True, det_limit_side_len=1280, use_angle_cls=False, lang='ch', det_model_dir='/root/.paddleocr/paddleocr_det_infer', rec_model_dir='/root/.paddleocr/paddleocr_rec_infer', cls_model_dir='/root/.paddleocr/paddleocr_cls_infer')
    
    app.config[SERVICE_TYPE] = service_type
    app.config[MODEL_PADDLE_OCR] = paddle_ocr
    app.config[MODEL_OCR_DETERTION] = ocr_detection
    app.config[MODEL_OCR_RECOGNITION] = ocr_recognition
    app.config[MODEL_TABLE_RECOGNITION] = table_recognition
    app.config[MODEL_IMAGE_ORIENTATION_PREDICTOR] = image_orientation_predictor
    print("initialize app end, cost: {:.3f} seconds".format(time.time() - start_time))

def create_app():
    app = Flask(__name__)

    @app.route("/predict/system", methods=['POST'])
    def ocr_predict():
        # 获取模型等配置数据
        service_type = current_app.config[SERVICE_TYPE]
        paddle_ocr = current_app.config[MODEL_PADDLE_OCR]
        ocr_detection = current_app.config[MODEL_OCR_DETERTION]
        ocr_recognition = current_app.config[MODEL_OCR_RECOGNITION]
        table_recognition = current_app.config[MODEL_TABLE_RECOGNITION]
        image_orientation_predictor = current_app.config[MODEL_IMAGE_ORIENTATION_PREDICTOR]
        
        start_time = time.time()
        # 获取请求体中的数据
        data = request.get_data() 
        json_data = json.loads(data)

        # 1. 加载图片
        images, load_images_cost = load_images(json_data)
        if images is None:
            return { 'status': -1, 'message': '参数异常：图片信息(images/urls)不能为空' }

        # 2. 图片预处理
        pre_images, pre_process_cost = pre_process_images(images, json_data)

        # 3. 方向纠正
        angle_process_images, angle_process_res, angle_cost = angle_images(pre_images, json_data, image_orientation_predictor)

        # 4. 文本信息提取
        ocr_res_future = executor.submit(read_images_ocr_words, angle_process_images, json_data, paddle_ocr, ocr_detection, ocr_recognition)

        # 5. 表格信息提取
        table_res_future = executor.submit(read_images_tables, angle_process_images, json_data, table_recognition)

        # 6. 条形码信息提取
        barcode_res_future = executor.submit(read_images_barcodes, angle_process_images, json_data)

        # 7. 主动释放显存
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        if hasattr(paddle.device.cuda, 'empty_cache'):
            paddle.device.cuda.empty_cache()

        # 等待结果 合并结果返回
        ocr_res, words_cost = ocr_res_future.result()
        table_res, table_cost = table_res_future.result()
        barcode_res, barcode_cost = barcode_res_future.result()
        res = merge_res(angle_process_res, ocr_res, table_res, barcode_res)
        
        total_cost = time.time() - start_time
        print("[OCR Request End] RT: {:.3f} seconds, Load Images Cost: {:.3f}, Pre Process Cost: {:.3f} Angle Cost: {:.3f}, Words Cost: {:.3f}, Table Cost: {:.3f}, BarCode Cost: {:.3f}".format(total_cost, load_images_cost, pre_process_cost, angle_cost, words_cost, table_cost, barcode_cost))
        return res

    @app.errorhandler(Exception)
    def handle_exception(e):
        # 打印请求的 URL
        print(f'Error in request to: {request.url}')
        # 打印请求的方法
        print(f'Request method: {request.method}')
         # 还可以记录其他信息，如远程地址
        print(f'Remote address: {request.remote_addr}')
        # 打印请求头
        print(f'Request headers: {request.headers}')
        # 打印异常堆栈
        traceback.print_exc()
        # 继续正常的错误处理
        return str(e), 500
    
    return app

def load_images(json_data):
    start_time = time.time()
    images = json_data.get('images', None)
    urls = json_data.get('urls', None)
    if images is None and urls is None:
        return None

    # 支持图片链接识别
    if images is None:
        images = []
        for url in urls:
            response = requests.get(url)
            images.append(base64.b64encode(response.content).decode())
    cost = time.time() - start_time
    return images, cost

def merge_res(base_res, ocr_res, table_res, barcode_res):
    # 合并结果
    for index, result in enumerate(base_res):
        result['words'] = ocr_res[index]
        result['barcodes'] = barcode_res[index]
        result['tables'] = table_res[index]
    return {
        'status': '000',
        'message': None,
        'results': base_res
    }

app = create_app()
if __name__ == "__main__":
    app.run(host='0.0.0.0')